import pytest
import torch
import numpy as np
import sys
import os

# 添加项目根目录到Python路径
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from src.operators.asecondc import AsecondC

def test_asecondc_output_shape():
    """测试AsecondC算子的输出形状是否正确"""
    batch_size = 2
    in_channels = 3
    height = 32
    width = 32
    out_channels = 16
    kernel_size = 3
    
    # 创建输入
    x = torch.randn(batch_size, in_channels, height, width)
    
    # 创建算子实例
    asecondc = AsecondC(
        in_channels=in_channels,
        out_channels=out_channels,
        kernel_size=kernel_size,
        stride=1,
        padding=1
    )
    
    # 计算输出
    output = asecondc(x)
    
    # 验证输出形状
    expected_height = height  # 因为padding=1，所以高度保持不变
    expected_width = width    # 因为padding=1，所以宽度保持不变
    expected_shape = (batch_size, out_channels, expected_height, expected_width)
    
    assert output.shape == expected_shape, \
        f"输出形状 {output.shape} 不符合预期 {expected_shape}"

def test_asecondc_backward():
    """测试AsecondC算子的反向传播"""
    batch_size = 2
    in_channels = 3
    height = 32
    width = 32
    out_channels = 16
    kernel_size = 3
    
    # 创建输入，需要设置requires_grad=True以支持反向传播
    x = torch.randn(batch_size, in_channels, height, width, requires_grad=True)
    
    # 创建算子实例
    asecondc = AsecondC(
        in_channels=in_channels,
        out_channels=out_channels,
        kernel_size=kernel_size,
        stride=1,
        padding=1
    )
    
    # 前向传播
    output = asecondc(x)
    
    # 创建一个虚拟的损失
    loss = output.sum()
    
    # 尝试反向传播
    try:
        loss.backward()
    except Exception as e:
        pytest.fail(f"反向传播失败: {str(e)}")
    
    # 验证梯度不为None
    assert x.grad is not None, "输入的梯度为None"
    assert asecondc.weight.grad is not None, "权重的梯度为None"
    assert asecondc.bias.grad is not None, "偏置的梯度为None"

def test_asecondc_device_compatibility():
    """测试AsecondC算子在不同设备上的兼容性"""
    if not torch.cuda.is_available():
        pytest.skip("CUDA不可用，跳过GPU测试")
    
    batch_size = 2
    in_channels = 3
    height = 32
    width = 32
    out_channels = 16
    kernel_size = 3
    
    # 创建算子实例
    asecondc = AsecondC(
        in_channels=in_channels,
        out_channels=out_channels,
        kernel_size=kernel_size,
        stride=1,
        padding=1
    )
    
    # 测试CPU
    x_cpu = torch.randn(batch_size, in_channels, height, width)
    output_cpu = asecondc(x_cpu)
    assert output_cpu.device.type == "cpu"
    
    # 测试GPU
    asecondc = asecondc.cuda()
    x_gpu = x_cpu.cuda()
    output_gpu = asecondc(x_gpu)
    assert output_gpu.device.type == "cuda"

def test_asecondc_numerical_values():
    """测试AsecondC算子的数值计算是否正确"""
    # 创建一个简单的输入用例
    x = torch.ones(1, 1, 4, 4)
    
    # 创建一个简单的AsecondC实例
    asecondc = AsecondC(
        in_channels=1,
        out_channels=1,
        kernel_size=3,
        stride=1,
        padding=1
    )
    
    # 设置固定的权重和偏置以便进行确定性测试
    with torch.no_grad():
        asecondc.weight.fill_(1.0)
        asecondc.bias.fill_(0.0)
    
    # 计算输出
    output = asecondc(x)
    
    # 对于给定的输入和权重，我们可以手动计算期望的输出值
    # 在这个例子中，中心位置的输出应该是9（3x3的kernel全是1，输入也全是1）
    expected_center_value = 9.0
    
    # 验证中心位置的值
    center_value = output[0, 0, 2, 2].item()
    assert abs(center_value - expected_center_value) < 1e-6, \
        f"中心位置的值 {center_value} 不等于期望值 {expected_center_value}"

if __name__ == "__main__":
    pytest.main([__file__]) 